#include "cmainwindow.h"
#include <buola/scene/cperspectivecamera.h>
#include <buola/scene/controllers/ccameracontroller.h>
#include <buola/scene/transform/cmattransform.h>
#include <buola/scene/geodes/ccube.h>
#include <buola/cv/io/ckinect.h>
#include <buola/cv/features/cfeatureset.h>
#include <buola/cv/features/cfeaturematches.h>
#include <buola/cv/features/cmatcher.h>
#include <buola/cv/geometry/transform.h>
#include <buola/cv/registration/icp.h>
#include <buola/widgets/ctoolbar.h>
#include <buola/gui/csimpleapp.h>
#include <buola/gui/ctoolkit.h>
#include <buola/image/cgc.h>
#include <buola/threads/thread.h>

namespace buola { namespace reconstructor {

CMainWindow::CMainWindow()
    :   mTimer(100_ms,true)
{
    mTimer.sTimer.Connect(&CMainWindow::OnTimer,this);
    
    mToolBar=new gui::CToolBar();
    Attach(mToolBar,gui::ATTACH_NORTH,gui::ATTACH_TOOLBAR);

    gui::PMenu lMenu=new gui::CMenu;
    
    lMenu->Add(gui::new_menu_item(L"start camera",ICON_STOCK_UP))->
              eTriggered.Connect(&CMainWindow::StartCamera,this);
    lMenu->Add(gui::new_menu_item(L"start rec",ICON_STOCK_ARROW_RIGHT))->
              eTriggered.Connect(&CMainWindow::StartRec,this);
    lMenu->Add(gui::new_menu_item(L"start load",ICON_STOCK_ARROW_RIGHT))->
              eTriggered.Connect(&CMainWindow::StartLoad,this);
    
    mToolBar->SetMenu(lMenu);

    SetupScene();
    SetupViews();
}

CMainWindow::~CMainWindow()
{
}

void CMainWindow::SetupScene()
{
    mScene=new scene::CScene;
    
    mCamera=new scene::CPerspectiveCamera;
    mCamera->SetFOV(2*atan2(240.0,525.0));
    mCamera->SetClipping(0.1,50.0);
    mCamera->LookAt({0,0,1},{0,0,0},{0,-1,0});
}

void CMainWindow::SetupViews()
{
    mSceneView=new scene::CSceneView;
    
    mSceneView->AddScene(mScene);
    mSceneView->SetCamera(mCamera);
    mSceneView->AddController(new scene::CCameraController(mCamera));
    mSceneView->Resize({1280,960});
    mSceneView->SetTitle(L"current iteration");
    
    Add(mSceneView);
}

void CMainWindow::StartCamera()
{
    mKinect.reset(new cvi::CKinect);
    mKinect->sImages.Connect([this](const img::CImage_rgb8 &pRGB,const img::CImage_gray32f &pDepth){OnCapture(pRGB,pDepth);});
    mKinect->Start();
}

void CMainWindow::StartRec()
{
    in_thread([this]{RecThread();})();
}

void CMainWindow::StartLoad()
{
    in_thread([this]{LoadThread();})();
}


void CMainWindow::OnCapture(const img::CImage_rgb8 &pRGB,const img::CImage_gray32f &pDepth)
{
   mCaptureMutex.lock();
   mLastRGB=share(pRGB);
   mLastDepth=share(pDepth);
   mCaptureMutex.unlock();
}

void CMainWindow::RecThread()
{
    try
    {
        while(true)
        {
            img::CImage_rgb8 lRGB;
            img::CImage_gray32f lDepth;

            mCaptureMutex.lock();
            if(!mLastRGB)
            {
                mCaptureMutex.unlock();
                sleep(10_ms);
                continue;
            }
            lRGB=std::move(mLastRGB);
            lDepth=std::move(mLastDepth);
            mCaptureMutex.unlock();

            ProcessCloud(std::move(lRGB),std::move(lDepth));
        }
    }
    catch(std::exception &pE)
    {
        msg_info() << "caught exception in thread: " << pE.what() << "\n";
    }
}

void CMainWindow::LoadThread()
{
    try
    {
        while(true)
        {
            for(int i=0;;i++)
            {
                img::CImage_rgb8 lRGB;
                img::CImage_gray32f lDepth;
                img::load_rgbd(lRGB,lDepth,io::file("/home/xavi/clouds/purplething/purple"+to_string(i)+".pcd"));
                ProcessCloud(std::move(lRGB),std::move(lDepth));
            }
        }
    }
    catch(std::exception &pE)
    {
        msg_info() << "caught exception in thread: " << pE.what() << "\n";
    }
}

void CMainWindow::ProcessCloud(img::CImage_rgb8&& pRGB,img::CImage_gray32f &&pDepth)
{
    img::CCameraParameters lParams;
    lParams.SetMatrix(CPoint_d(525,525),CPoint_d(320,240));
    img::CImage_xyz32f lXYZ=cvi::make_pos_image(pDepth,lParams);
    
    img::CImage_gray8 lGray(pRGB.Size());
    img::convert_pixels(pRGB,lGray);

    img::CImage_gray8 lMask(pRGB.Size());
        
    for(int x=0;x<640;x++)
    {
        for(int y=0;y<480;y++)
        {
            lMask(x,y)=is_nan(pDepth(x,y))?0:255;
        }
    }
    
    cvi::CFeatureSet lFeatures=mSURF.Detect(lGray,lMask);

    mat::CMat34d lTrans=mat::identity_3d();

    cvi::CCloud_rgb8 lDownSampled;
    cvi::CCloud_rgb8 lPointCloud;
    
    if(!mKeyFrames.empty())
    {
        SKeyFrame &lLastKey=*mKeyFrames.back();
        
        cvi::CFeatureMatches lMatches=cvi::bf_match(lLastKey.mFeatures,lFeatures);
        lMatches.Sort();

        std::vector<mat::CVec3d> lLeft,lRight;
        
        for(int j=0;j<lMatches.size();j++)
        {
            const cvi::CKeyPoint &lKP1=lMatches.KeyPoint1(j);
            const cvi::CKeyPoint &lKP2=lMatches.KeyPoint2(j);
            mat::CVec3d lP1=lLastKey.mXYZ(lKP1.Point_i());
            mat::CVec3d lP2=lXYZ(lKP2.Point_i());
            if(is_nan(lP1.z())||is_nan(lP2.z())) continue;
            double lCoeff=lP1.z()*(double)lKP1.Size()/lP2.z()/(double)lKP2.Size();
            if(lCoeff>1.1||lCoeff<0.9) continue;
            lLeft.push_back(lP1);
            lRight.push_back(lP2);
        }
        
        mat::CMat34d lNewTrans=cvi::estimate_rigid_transform(lLeft,lRight,0.1,0.998);
        
        lTrans=lLastKey.mTrans*inverse(lNewTrans);
        
        double lDiff1=norm(lTrans*mat::CVec3d(0,0,0)-lLastKey.mTrans*mat::CVec3d(0,0,0));
        double lDiff2=norm(lTrans*mat::CVec3d(0,0,3)-lLastKey.mTrans*mat::CVec3d(0,0,3));
        
        if(lDiff1<0.5&&lDiff2<0.5) return;

        lPointCloud=cvi::make_cloud(lXYZ,pRGB);
        lDownSampled=cvi::downsample(lPointCloud,10);
        
        lNewTrans=cvi::icp(lDownSampled,lLastKey.mKDTree,inverse(lNewTrans));
        lTrans=lLastKey.mTrans*lNewTrans;
    }
    else
    {
        lPointCloud=cvi::make_cloud(lXYZ,pRGB);
        lDownSampled=cvi::downsample(lPointCloud,10);
    }
    msg_info() << "adding frame " << mKeyFrames.size() << "!!\n";

    SKeyFrame *lKeyFrame=new SKeyFrame;
    
    lKeyFrame->mRGB=std::move(pRGB);
    lKeyFrame->mXYZ=std::move(lXYZ);
    lKeyFrame->mFeatures=std::move(lFeatures);
    lKeyFrame->mDownSampled=std::move(lDownSampled);
    lKeyFrame->mTrans=lTrans;
    lKeyFrame->mKDTree.SetContent(lPointCloud);
    mKeyFrames.emplace_back(lKeyFrame);
    
    mPendingMutex.lock();
    mPendingKeyFrames.push_back(lKeyFrame);
    mPendingMutex.unlock();
}

void CMainWindow::OnTimer()
{
    while(true)
    {
        SKeyFrame *lKeyFrame=nullptr;
        
        mPendingMutex.lock();
        if(!mPendingKeyFrames.empty())
        {
            lKeyFrame=mPendingKeyFrames.back();
            mPendingKeyFrames.pop_back();
        }
        mPendingMutex.unlock();
        
        if(!lKeyFrame) break;
        
        scene::PPointCloud lSceneCloud=new scene::CPointCloud;
        lSceneCloud->SetCloud(cvi::CCloud_rgb8(lKeyFrame->mCloud));
        scene::PMatTransform lMatTransform=new scene::CMatTransform(lKeyFrame->mTrans);
        lSceneCloud->AttachTo(lMatTransform);
        scene::PCube lCube=new scene::CCube(0.1,0.1,0.1);
        lCube->AttachTo(lMatTransform);
    
        mScene->World()->AddChild(lMatTransform);
        mSceneView->Refresh();
    }
}

/*namespace reconstructor*/ } /*namespace buola*/ }

int main(int pNArgs,char **pArgs)
{
    using namespace buola;
    
    buola_init(pNArgs,pArgs);

    try
    {
        gui::get_toolkit()->GetOpenGLGC()->MakeCurrent();
        gui::CSimpleApp<reconstructor::CMainWindow> lApp;
        lApp.Run();
    }
    catch(std::exception &pE)
    {
        msg_error() << pE.what() << "\n";
    }

    return buola_finish();
}
